import numpy as np
from scipy.spatial.transform import Rotation as ScipyRotation


def rotation_matrix(axis, angle):
    """
    根据旋转轴和角度生成旋转矩阵（3x3）
    
    参数:
        axis: 旋转轴 [x, y, z]
        angle: 旋转角度（弧度）
        
    返回:
        3x3 旋转矩阵
    """
    axis = np.asarray(axis)
    axis = axis / np.linalg.norm(axis)  # 归一化
    
    # 使用罗德里格斯旋转公式
    cos_theta = np.cos(angle)
    sin_theta = np.sin(angle)
    one_minus_cos = 1.0 - cos_theta
    
    # 叉积矩阵
    cross_prod_mat = np.array([
        [0, -axis[2], axis[1]],
        [axis[2], 0, -axis[0]],
        [-axis[1], axis[0], 0]
    ])
    
    # 旋转矩阵
    rot_mat = cos_theta * np.eye(3) + one_minus_cos * np.outer(axis, axis) + sin_theta * cross_prod_mat
    
    return rot_mat

def print_rotation_info(rotation_matrix):
    """
    打印旋转矩阵的各种表示形式
    """
    # 从旋转矩阵创建 Rotation 对象
    r = ScipyRotation.from_matrix(rotation_matrix)
    
    # 1. 欧拉角表示 (zyx顺序，内旋，度)
    euler_angles_zyx_degrees = r.as_euler('zyx', degrees=True)
    
    # 2. 轴角表示 (轴向量和角度)
    # as_rotvec() 返回的向量方向是旋转轴，范数是旋转角度(弧度)
    rotvec = r.as_rotvec()
    angle_rad = float(np.linalg.norm(rotvec))
    if angle_rad > 0:  # 避免除以零
        axis = rotvec / angle_rad  # 单位向量表示旋转轴
    else:
        axis = np.zeros(3)
    angle_deg = float(np.degrees(angle_rad))
    
    # 3. 四元数表示
    quat = r.as_quat()  # 返回 (x, y, z, w) 格式
    
    # 将NumPy数组转换为普通列表
    euler_list = euler_angles_zyx_degrees.tolist()
    axis_list = axis.tolist()
    quat_list = quat.tolist()
    
    # print("-"*50)
    print(f"欧拉角 (zyx, 度): [{euler_list[0]:.4f}, {euler_list[1]:.4f}, {euler_list[2]:.4f}]")
    
    return {
        'euler_zyx_deg': euler_list,
        'axis': axis_list,
        'angle_deg': angle_deg,
        'angle_rad': angle_rad,
        'quaternion': quat_list
    }

def print_pose_info(pose_matrix):
    """
    打印姿态矩阵的信息，包括旋转和平移部分
    
    参数:
        pose_matrix: 4x4 齐次变换矩阵，表示姿态
        
    返回:
        包含旋转和平移信息的字典
    """
    # 提取旋转矩阵和平移向量
    rotation_matrix = pose_matrix[:3, :3]
    translation = pose_matrix[:3, 3]
    
    # 将NumPy数组转换为普通列表
    trans_list = translation.tolist()
    
    # 打印平移向量
    print(f"平移向量 (x, y, z): [{trans_list[0]:.4f}, {trans_list[1]:.4f}, {trans_list[2]:.4f}]")
    
    # 使用 print_rotation_info 打印旋转部分
    rotation_info = print_rotation_info(rotation_matrix)
    
    # 合并结果
    result = {
        'translation': trans_list,
        **rotation_info
    }
    
    return result

# 测试代码
if __name__ == "__main__":
    # 原始的齐次变换矩阵
    # T_to_fangnuo_left_wrist = np.array([[0, 0, 1, 0],
    #                                   [0, 1, 0, 0],
    #                                   [-1, 0, 0, 0],
    #                                   [0, 0, 0, 1]])
    T_robot_openxr = np.array([[1, 0, 0, 0],
                                    [0, 0, -1, 0],
                                    [0, 1, 0, 0],
                                    [0, 0, 0, 1]])
    # 提取旋转矩阵部分
    rotation_matrix = T_robot_openxr[:3, :3]
    
    # 打印旋转信息
    print_rotation_info(rotation_matrix)


def euler_to_matrix(euler_angles, sequence='zyx', degrees=True, position=None):
    """
    将欧拉角转换为4x4齐次变换矩阵
    
    参数:
        euler_angles: 欧拉角 [x, y, z] 或 [x, y, z, w]
        sequence: 欧拉角旋转顺序，例如 'xyz', 'zyx' 等
        degrees: 如果为True，则输入角度为度；如果为False，则为弧度
        position: 可选，平移向量 [x, y, z]，默认为 [0, 0, 0]
        
    返回:
        4x4 齐次变换矩阵
    """
    if position is None:
        position = [0, 0, 0]
    
    # 创建旋转对象
    rotation = ScipyRotation.from_euler(sequence, euler_angles[:3], degrees=degrees)
    
    # 获取3x3旋转矩阵
    rot_matrix = rotation.as_matrix()
    
    # 创建4x4齐次变换矩阵
    transform = np.eye(4)
    transform[:3, :3] = rot_matrix
    transform[:3, 3] = position
    
    return transform


# 测试代码
if __name__ == "__main__":
    # 测试欧拉角转矩阵
    print("\n测试欧拉角转矩阵:")
    euler_angles = [0, 95, 0]  # 绕y轴旋转95度
    transform = euler_to_matrix(euler_angles, degrees=True, position=[1, 2, 3])
    print(f"欧拉角 {euler_angles} 转换为变换矩阵:\n{transform}")
    
    # 验证转换的正确性
    print("\n验证转换结果:")
    print_rotation_info(transform[:3, :3])
    print(f"平移向量: {transform[:3, 3]}")

    print("-------------------------")
    # T = np.array([[-0.34202014, 0.0, 0.93969262, 0.0],
    #                                [0.0, 1.0, 0.0, 0.0],
    #                                [-0.93969262, 0.0, -0.34202014, 0.0],
    #                                [0.0, 0.0, 0.0, 1.0]])
    # print_rotation_info(T[:3, :3])